//
// Created by mabinjie on 2024/11/1.
//

#ifndef PACKET_HANDLER_HPP
#define PACKET_HANDLER_HPP
#include <numeric>
#include <set>

#include "protocol/tools/packet_factory.hpp"

namespace network {

class PacketHandler;
typedef std::shared_ptr<PacketHandler> handler_t;
typedef std::vector<std::shared_ptr<common::Serialized>> serialize_t;

/**
 * 该类用于组包和分包
 * 组包：将二进制数据反序列化为结构化的网络包，对分包数据会进行异步组包，组包完成时进行回调
 * 分包：对结构化的网络包进行分包和序列化
 * 该类使用工厂模式，一个连接会话使用一个处理类
 */
class PacketHandler {
public:
    PacketHandler() = default;
    explicit PacketHandler(const std::string &id, const std::function<void(protocol::Packet)> &callback) : callback_(callback) {
        packet_factory = protocol::PacketFactory::factory(id);
    }


    /**
     * 用于获取包处理类
     * @param id 工厂id，一般传远端地址
     * @param callback 回调函数，会在组包完成时调用
     * @return 处理类
     */
    static handler_t create(const std::string &id,
                            const std::function<void(protocol::Packet)>& callback) {
        std::call_once(flag, []{
            handlers = new std::unordered_map<std::string, handler_t>();
        });

        if (handlers->contains(id))
            return (*handlers)[id];

        (*handlers)[id] = std::make_shared<PacketHandler>(id, callback);
        return (*handlers)[id];
    }

    /**
     * 序列化操作
     * @param serializable 可序列化对象
     * @return 序列化数据
     * TODO writer大小可配置
     */
    static std::shared_ptr<common::Serialized> serialize(common::Serializable *serializable) {
        writer_t writer(2048);
        serializable->serialize(writer);
        return writer.flushAndGet();
    }

    /**
     * 对传入的包进行分包处理并序列化
     * @param packet 要进行分包的的对象
     * @param max_packet_size 分包阈值，包大小小于阈值时抛出异常
     * @return 分包后的序列化对象
     */
    [[nodiscard]] serialize_t segment(const protocol::jtt808_packet &packet, const size_t max_packet_size = 1024) const {
        if (size(packet) < max_packet_size)
            throw std::invalid_argument("Packet is too small");

        const auto max_body_size = max_packet_size - 21 - 3;
        const auto complete_body = packet.body_;
        const auto total_segment = static_cast<int>(std::ceil(static_cast<float>(size(*complete_body)) / max_body_size));
        const auto serialized_body = serialize(complete_body.get());
        std::vector<protocol::jtt808_body_ptr> bodies;
        for (int i = 0; i < total_segment; ++i) {
            auto part_body = serialized_body->sub(i*max_body_size, max_body_size);
            auto partial = std::make_shared<protocol::PartialBody>(*part_body, complete_body->message_id_);
            bodies.push_back(partial);
        }
        auto packets = packet_factory->packets(bodies, packet.header_.attribute_.is_encrypted);
        serialize_t result(packets.size());
        std::ranges::transform(packets.begin(), packets.end(), result.begin(),
                               [](const protocol::Packet &p) { return p->serialize(); });
        return result;
    }

    /**
     * 从序列化后的packet中局部反序列化消息头
     * @param serialized 序列化后的packet
     * @return 消息头
     */
    static std::shared_ptr<protocol::jtt808_header> read_header(const std::shared_ptr<common::Serialized> &serialized) {
        const auto first_header = serialized->sub(BEGIN_OF_HEADER, LENGTH_OF_HEADER);
        reader_t reader(first_header->data, first_header->size, protocol::reverse_escaper);
        auto header = std::make_shared<protocol::jtt808_header>();
        header->deserialize(reader);
        return header;
    }

    /**
     * 某一消息的所有分包到达后进行组包
     * @param serialize 完整的一组序列化数据
     * @return 反序列化后的包
     */
    [[nodiscard]] protocol::Packet reassembly(const serialize_t &serialize) const {
        const auto header = read_header(serialize[0]);
        auto body_parts = std::vector<std::shared_ptr<common::Serialized>>(serialize.size());
        for (int i = 0; i < serialize.size(); ++i) {
            const auto body_part = get_body_part(*serialize[i]);
            body_parts[i] = body_part;
        }
        const auto one = common::Serialized::all_in_one(body_parts);
        header->attribute_.is_segment = false;
        const auto body = protocol::MessageFactory::create(*header);
        reader_t r(one->data, one->size, protocol::reverse_escaper);
        body->deserialize(r);
        auto attribute = header->attribute_;
        attribute.message_length = size(*body);
        return packet_factory->packet(body, attribute);
    }

    /**
     * 检查分包数据是否完整
     * @param to_check 要检查的二进制数据
     * @return 是否完整
     * TODO 优化检查逻辑提高性能
     */
    static bool check_for_integrity(const serialize_t &to_check) {
        if (to_check.empty()) return false;
        std::set<uint16_t> packet_indexes;
        for (const auto& serialized : to_check) {
            const auto header = read_header(serialized);
            packet_indexes.insert(header->segmentation_.packet_index);
        }
        if (const auto h = read_header(to_check[0]); packet_indexes.size() == h->segmentation_.packet_total)
            return true;
        return false;
    }

    /**
     * 异步接收收到的数据，如果不是分包消息则直接组包调用回调函数，如果是分包消息，则在所有分包到达后调用回调函数
     * @param serialized 本次收到的数据
     */
    void async_handle(const std::shared_ptr<common::Serialized>& serialized) {
        const auto header = read_header(serialized);
        if (!header->attribute_.is_segment) {
            const auto body = protocol::MessageFactory::create(*header);
            const auto body_part = get_body_part(*serialized);
            reader_t reader(body_part->data, body_part->size, protocol::reverse_escaper);
            body->deserialize(reader);
            callback_(std::make_shared<protocol::jtt808_packet>(*header, body));
            return;
        }
        if (to_reassembly.contains(header->message_id_)) {
            auto &reassembly_list = to_reassembly[header->message_id_];
            reassembly_list.push_back(serialized);
            if (check_for_integrity(reassembly_list)) {
                callback_(reassembly(reassembly_list));
                to_reassembly.erase(header->message_id_);
            }
        } else {
            to_reassembly[header->message_id_] = serialize_t({serialized});
        }
    }

    /**
     * 进一步组包，主要用于主发报文时延续消息序列号
     * @param body 消息体
     * @param attribute 消息头属性
     * @param segmentation 分包属性
     * @return 组包结果
     */
    protocol::Packet packet(const protocol::jtt808_body_ptr& body, const protocol::MessageAttribute &attribute,
                            const protocol::segmentation &segmentation = {}) {
        return packet_factory->packet(body, attribute, segmentation);
    }

    // 由于一个连接绑定一个handler，因此在此处记录本条连接的终端id和imei
    int terminal_id = -1;
    std::vector<uint8_t> terminal_imei;

private:
    static inline std::unordered_map<std::string, handler_t> *handlers = nullptr;
    static inline std::once_flag flag;

    std::unordered_map<uint32_t, serialize_t> to_reassembly;
    protocol::packet_factory_t packet_factory;
    std::function<void(protocol::Packet)> callback_;

    /**
     * 从完整的系列化数据中提取消息体部分
     * @param all 完整的序列化数据
     * @return 消息体序列化数据
     */
    static std::shared_ptr<common::Serialized> get_body_part(const common::Serialized &all) {
        return all.sub(BEGIN_OF_MESSAGE, LENGTH_OF_MESSAGE(all.size));
    }

};

}

#endif //PACKET_HANDLER_HPP
